import tensorflow as tf
# 加载MNIST数据集
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()
# 数据预处理（归一化 + One-Hot编码）
def preprocess_data(images, labels):
    # 归一化到 [0,1]
    images = images.astype('float32') / 255.0
    # 转换维度 (60000, 28, 28) -> (60000, 784)
    images = images.reshape((-1, 28*28))
    # One-Hot编码
    labels = tf.keras.utils.to_categorical(labels, 10)
    return images, labels
# 处理训练集和测试集
train_images, train_labels = preprocess_data(train_images, train_labels)
test_images, test_labels = preprocess_data(test_images, test_labels)
# 验证集（取训练集前5000个样本）
val_images, val_labels = train_images[:5000], train_labels[:5000]
train_images, train_labels = train_images[5000:], train_labels[5000:]
# 打印数据形状
print("训练集:", train_images.shape, train_labels.shape)  # (55000, 784) (55000, 10)
print("验证集:", val_images.shape, val_labels.shape)      # (5000, 784) (5000, 10)
print("测试集:", test_images.shape, test_labels.shape)    # (10000, 784) (10000, 10)